from pydantic import BaseModel, Field
from typing import Dict, List, Any, Optional, Union
from datetime import datetime
from enum import Enum

# ============================================================================
# 数据模型定义文件
# 定义了智能助手系统中使用的所有数据结构和类型
# 包括消息类型、用户信息、会话上下文、工作流状态等
# ============================================================================

class MessageType(str, Enum):
    """消息类型枚举
    
    定义系统中支持的所有消息类型，用于WebSocket通信和消息路由
    """
    USER = "user"                    # 用户输入消息
    ASSISTANT = "assistant"          # 助手回复消息
    CHART = "chart"                  # 图表数据消息
    FORM = "form"                    # 表单数据消息
    CONFIRMATION = "confirmation"    # 确认请求消息
    STATUS = "status"                # 状态更新消息
    ERROR = "error"                  # 错误消息
    HEARTBEAT = "heartbeat"          # 心跳消息

class WorkflowStatus(str, Enum):
    """工作流状态枚举
    
    定义工作流的各种执行状态，用于跟踪多Agent协作的进度
    """
    RUNNING = "running"                          # 正在运行
    WAITING_CONFIRMATION = "waiting_confirmation" # 等待用户确认
    COMPLETED = "completed"                      # 已完成
    FAILED = "failed"                            # 执行失败
    CANCELLED = "cancelled"                      # 已取消

class UserInfo(BaseModel):
    """用户信息模型
    
    存储用户的基本信息，用于身份验证和权限控制
    """
    user_id: str        # 用户唯一标识符
    name: str           # 用户姓名
    department: str     # 所属部门
    position: str       # 职位
    email: str          # 邮箱地址
    tenant_id: str = "tenant001"  # 租户ID，默认为tenant001

class TenantInfo(BaseModel):
    """租户信息模型
    
    多租户系统中的租户配置信息，包含权限和功能限制
    """
    tenant_id: str          # 租户唯一标识符
    tenant_name: str        # 租户名称
    permissions: List[str]  # 权限列表
    max_users: int          # 最大用户数限制
    features: List[str]     # 可用功能列表

class SessionContext(BaseModel):
    """会话上下文模型
    
    维护用户会话的完整上下文信息，包括用户信息、租户信息和会话状态
    """
    session_id: str                                          # 会话唯一标识符
    user_info: UserInfo                                      # 用户信息
    tenant_info: TenantInfo                                  # 租户信息
    created_at: datetime = Field(default_factory=datetime.now)     # 会话创建时间
    last_activity: datetime = Field(default_factory=datetime.now)  # 最后活动时间
    workflow_state: Optional[Dict[str, Any]] = None         # 工作流状态数据
    skill_selected: Optional[str] = None                    # 当前选择的技能
    enable_network: bool = False                            # 是否启用网络功能

class BaseMessage(BaseModel):
    """基础消息模型
    
    所有消息类型的基础结构，包含消息的基本元数据
    """
    message_id: str                                          # 消息唯一标识符
    session_id: str                                          # 所属会话ID
    message_type: MessageType                                # 消息类型
    timestamp: datetime = Field(default_factory=datetime.now)     # 消息时间戳
    content: Dict[str, Any]                                  # 消息内容

class UserMessage(BaseModel):
    """用户消息模型
    
    用户发送的消息结构，包含查询内容和相关配置
    """
    message: str                        # 用户输入的消息内容
    skill_selected: Optional[str] = None  # 用户选择的技能类型
    enable_network: bool = False        # 是否启用网络搜索功能

class AssistantMessage(BaseModel):
    """助手消息模型
    
    助手回复的消息结构，支持流式输出和工作流关联
    """
    content: str                        # 助手回复的内容
    is_streaming: bool = False          # 是否为流式输出
    is_final: bool = True               # 是否为最终消息
    workflow_id: Optional[str] = None   # 关联的工作流ID

class ChartData(BaseModel):
    """图表数据模型
    
    用于数据可视化的图表配置和数据结构
    支持多种图表类型：line(折线图), bar(柱状图), pie(饼图), table(表格)
    """
    chart_type: str                                    # 图表类型：line, bar, pie, table
    title: str                                         # 图表标题
    data: List[Dict[str, Any]]                        # 图表数据
    config: Dict[str, Any] = Field(default_factory=dict)    # 图表配置参数
    analysis_text: Optional[str] = None               # 图表分析文本
    is_streaming: bool = False                        # 是否为流式输出
    is_final: bool = True                             # 是否为最终数据

class FormField(BaseModel):
    """表单字段模型
    
    定义表单中单个字段的结构和属性
    支持多种字段类型：text, select, date, number, textarea
    """
    name: str                           # 字段名称（用于数据绑定）
    label: str                          # 字段显示标签
    type: str                           # 字段类型：text, select, date, number, textarea
    required: bool = True               # 是否为必填字段
    options: Optional[List[str]] = None # 选择框的选项列表（仅select类型使用）
    value: Optional[Any] = None         # 字段当前值
    placeholder: Optional[str] = None   # 占位符文本

class FormData(BaseModel):
    """表单数据模型
    
    完整的表单结构，包含所有字段和元数据
    """
    form_id: str                        # 表单唯一标识符
    title: str                          # 表单标题
    description: str                    # 表单描述
    fields: List[FormField]             # 表单字段列表
    entity_data: Dict[str, Any]         # 完整的实体数据（包含业务逻辑相关数据）
    submit_endpoint: Optional[str] = None  # 表单提交的API端点
    status: str = "draft"               # 表单状态：draft(草稿), confirmed(已确认), submitted(已提交)

class ConfirmationRequest(BaseModel):
    """确认请求模型
    
    用于多Agent协作中的人工确认机制
    当系统需要用户确认某个操作时使用
    """
    confirmation_id: str                # 确认请求的唯一标识符
    message: str                        # 确认提示消息
    options: List[str]                  # 可选的确认选项
    timeout: int = 300                  # 超时时间（秒），默认5分钟
    allow_custom: bool = True           # 是否允许用户自定义回复
    workflow_id: Optional[str] = None   # 关联的工作流ID

class ConfirmationResponse(BaseModel):
    """确认响应模型
    
    用户对确认请求的响应结构
    """
    confirmation_id: str                                    # 对应的确认请求ID
    response: str                                           # 用户的响应内容
    selected_options: List[str] = Field(default_factory=list)     # 用户选择的选项
    is_custom: bool = False                                 # 是否为自定义回复

class StatusUpdate(BaseModel):
    """状态更新模型
    
    用于通知工作流执行状态的变化
    """
    status: WorkflowStatus                  # 工作流状态
    message: str                            # 状态描述消息
    progress: Optional[int] = None          # 进度百分比（0-100）
    workflow_id: Optional[str] = None       # 关联的工作流ID
    details: Optional[Dict[str, Any]] = None  # 详细信息

class ErrorMessage(BaseModel):
    """错误消息模型
    
    系统错误信息的标准结构
    """
    error_code: str                         # 错误代码
    error_message: str                      # 错误描述
    details: Optional[Dict[str, Any]] = None  # 错误详细信息
    recoverable: bool = False               # 是否可恢复的错误

class HeartbeatMessage(BaseModel):
    """心跳消息模型
    
    用于维持WebSocket连接活跃状态
    """
    timestamp: datetime = Field(default_factory=datetime.now)  # 心跳时间戳
    status: str = "alive"                   # 连接状态

# ============================================================================
# Agent相关模型
# 定义Agent配置和工作流状态管理相关的数据结构
# ============================================================================

class AgentConfig(BaseModel):
    """Agent配置模型
    
    定义单个Agent的配置信息，包括能力、数据源和优先级
    """
    name: str                                       # Agent名称
    description: str                                # Agent功能描述
    keywords: List[str]                             # 关键词列表（用于意图匹配）
    data_sources: List[str]                         # 可访问的数据源列表
    output_types: List[str]                         # 支持的输出类型
    priority: int                                   # 优先级（数值越高优先级越高）
    examples: List[str]                             # 使用示例
    config: Dict[str, Any] = Field(default_factory=dict)  # 额外配置参数

class WorkflowState(BaseModel):
    """工作流状态模型
    
    跟踪多Agent协作工作流的执行状态和上下文数据
    """
    workflow_id: str                                # 工作流唯一标识符
    session_id: str                                 # 所属会话ID
    workflow_type: str                              # 工作流类型
    current_step: str                               # 当前执行步骤
    context_data: Dict[str, Any]                    # 上下文数据
    confirmation_data: Optional[Dict[str, Any]] = None  # 确认相关数据
    next_action: Optional[str] = None               # 下一步操作
    status: WorkflowStatus                          # 工作流状态
    created_at: datetime = Field(default_factory=datetime.now)   # 创建时间
    updated_at: datetime = Field(default_factory=datetime.now)   # 更新时间
    expires_at: Optional[datetime] = None           # 过期时间

# ============================================================================
# 数据源相关模型
# 定义数据查询和结果的标准结构
# ============================================================================

class DataSourceQuery(BaseModel):
    """数据源查询模型
    
    统一的数据源查询接口，支持多种数据源类型
    """
    source_type: str                                # 数据源类型：user_list, department_list, meeting_rooms, production_data, knowledge_base
    query_params: Dict[str, Any] = Field(default_factory=dict)  # 查询参数
    filters: Dict[str, Any] = Field(default_factory=dict)       # 过滤条件

class DataSourceResult(BaseModel):
    """数据源查询结果模型
    
    标准化的数据查询结果结构
    """
    source_type: str                                # 数据源类型
    data: List[Dict[str, Any]]                      # 查询结果数据
    total_count: int                                # 总记录数
    query_time: datetime = Field(default_factory=datetime.now)   # 查询时间
    metadata: Dict[str, Any] = Field(default_factory=dict)       # 元数据信息

# ============================================================================
# WebSocket消息包装
# 定义WebSocket通信的统一消息格式
# ============================================================================

class WebSocketMessage(BaseModel):
    """WebSocket消息包装模型
    
    WebSocket通信的统一消息格式，支持所有类型的消息数据
    """
    type: MessageType                               # 消息类型
    data: Union[                                    # 消息数据（支持多种类型）
        UserMessage,        # 用户消息
        AssistantMessage,   # 助手消息
        ChartData,          # 图表数据
        FormData,           # 表单数据
        ConfirmationRequest,  # 确认请求
        ConfirmationResponse, # 确认响应
        StatusUpdate,       # 状态更新
        ErrorMessage,       # 错误消息
        HeartbeatMessage    # 心跳消息
    ]
    session_id: str                                 # 会话ID
    message_id: str                                 # 消息ID
    timestamp: datetime = Field(default_factory=datetime.now)  # 消息时间戳